package ai.mspbots.sql.controller;

import ai.mspbots.sql.datasource.DataSourceContextHolder;
import ai.mspbots.sql.dto.DdlDTO;
import ai.mspbots.sql.dto.QueryDTO;
import ai.mspbots.sql.repository.SqlRepository;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestHeader;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.bind.annotation.RestController;

import java.net.URLDecoder;
import java.nio.charset.Charset;
import java.util.List;

@RestController
@RequestMapping("/ai")
@RequiredArgsConstructor(onConstructor = @__(@Autowired))
public class AiController {

    private final SqlRepository sqlRepository;

    @PostMapping("/sql/generate")
    @ResponseBody
    public List<DdlDTO> generate(@RequestBody String sql) {
        String sqlQuery = URLDecoder.decode(sql, Charset.defaultCharset());
        sqlQuery = sqlQuery.replaceAll("\\{tenant_code}", String.valueOf(1285403951449878530L));
        return sqlRepository.generateSql(sqlQuery);
    }

    @PostMapping("/sql/explain")
    @ResponseBody
    public List<QueryDTO> explainSql(@RequestBody String sql, @RequestHeader(value = "datasource", required = false, defaultValue = "msp") String datasource) {
        DynamicDataSourceContextHolder.push(datasource);
        String sqlQuery = URLDecoder.decode(sql, Charset.defaultCharset());
        sqlQuery = sqlQuery.replaceAll("\\{tenant_code}", String.valueOf(1285403951449878530L));
        List<QueryDTO> queryDTOList = sqlRepository.explainSql(sqlQuery);
        DynamicDataSourceContextHolder.clear();
        return queryDTOList;
    }
}
